import matplotlib.pyplot  as pyplot
import pylab as pl
import numpy as np
import math 

R = 8.3144621   # R is the gas constant [R]=[joule/(mol*kelvin)].

#----------------------------------------------------#
#                 Politropic Density                 #
def rho_e(gamma, T, rhob):
    '''
    This function return the politropic density.

    INPUT:
    > gamma: Index of politropic model;
    > T: Temperature array.
    > rhob: Density at bottom

    OUTPUT:
    > rho: Politropic density.
    '''
    T_aux = T**gamma
    rho = (rhob / T_aux[0]) * T_aux
    return rho                  # kg/m**3

#----------------------------------------------------#
#                  Figure .png                       #

def pdf_plot(pdfname):
    """ 
    This function creates a figure in format .png
    
    INPUT:
    > pdfname: figure name;
    
    OUTPUT:
    > Creates a picture in the directory where you are working.
    """
    plt.savefig(pdfname + '.png')

#----------------------------------------------------#
#              Figure in Mollweide                   #

def moll_plot(Z, nlevels):
    """
    This function creates a figure in Mollweide projection.
    
    INPUTS:
    > Z: field to be plotted;
    > nlevels: Number of levels in the colormap.
    
    OUTPUT:
    > Creates a picture with Mollweide projection in the directory where you are working
    """
    cmap = cm.hot               		                      # Set colormap
    m = Basemap(projection='moll', lon_0 = 0)
    ra = np.linspace(-np.pi, np.pi, Z.shape[0])
    dec= np.linspace(-np.pi/2, np.pi/2, Z.shape[1])
    RAD = 180/np.pi
    X,Y = np.meshgrid(ra*RAD,dec*RAD) 
    levels = np.linspace(Z.min(),Z.max(), nlevels)
    return m.contourf(X, Y, Z, cmap = cmap, levels = levels, latlon = True)

#----------------------------------------------------#
#                  Mean Values                       #

def means(inputfile, nn):
    """ 
    This functions calculates the mean values for the filds in the inputfile (the output os the readxav function).
    
    INPUTS: 
    > inputfile: File to be read;
    > nn: number of slices in the calculus of average;
    
    OUTPUT:
    > mean_ : mean with nn slices.
    """
    inputs = np.load(inputfile)
    list1 = inputs['thew'], inputs['u'], inputs['v'], inputs['w'], inputs['ox'], inputs['oy'], inputs['oz'], inputs['u2'], inputs['v2'], inputs['w2'], inputs['ox2'], inputs['oy2'], inputs['oz2'], inputs['rwv'], inputs['rwu'], inputs['rvu'], inputs['bx'], inputs['by'], inputs['bz'], inputs['bx2'], inputs['by2'], inputs['bz2'], inputs['bzby'], inputs['bzbx'], inputs['bxby'], inputs['p'], inputs['the'] 
    nt = list1[0].shape[2]
    for i in len(list1):
        list1[i] = list[i][:, :, (nt - 1 - nn):(nt - 1)].sum(axis = 2)/nn 	
    mean_thew, mean_u, mean_v, mean_w, mean_ox, mean_oy, mean_oz, mean_u2, mean_v2, mean_w2, mean_ox2, mean_oy2, mean_oz2, mean_rwv, mean_rwu, mean_rvu, mean_bx, mean_by, mean_bz, mean_bx2, mean_by2, mean_bz2, mean_bzby, mean_bzbx, mean_bxby, mean_p, mean_the  = list1
    return mean_thew, mean_u, mean_v, mean_w, mean_ox, mean_oy, mean_oz, mean_u2, mean_v2, mean_w2, mean_ox2, mean_oy2, mean_oz2, mean_rwv, mean_rwu, mean_rvu, mean_bx, mean_by, mean_bz, mean_bx2, mean_by2, mean_bz2, mean_bzby, mean_bzbx, mean_bxby, mean_p, mean_the

#----------------------------------------------------#
#               Potencial Temperature                #

def PT(T, P, gamma):     
    '''
    This function calculates the potencial temperature using a c_p varying with the radius.
    
    INPUTS:
    > P: Pressure array;
    > r: Radius array;
    > gamma: Gamma array or scalar.

    OUTPUT:
    > theta: Potential temperature, in kelvin.
    '''
    theta = np.empty(P.size)
    theta = T * (P[0] / P)**(1.-1./gamma)
    return theta                

#----------------------------------------------------#
#                   Read BOX                         #

def rbox(inputfile, n, m, l, nt):
    '''
    This function is designed for read the BOX ('.fort11') written in order ['u', 'v', 'w', 't', 'p', 'bx', 'by', 'bz'], where this elements are matrices with order (n, m, l).  

    INPUTS: 
    > Inputfile: file to be read
    > n: number of steps in phi direction
    > m: number of steps in theta direction
    > l: number of steps in radius direction
    > nt: number of time steps
    
    OUTPUT:
    > Matrices in order (n, m, l, nt)
    '''
    element = np.zeros((n, m, l))                                       
    list1 = [element,]*8                                              
    list2 = list1                                            
    f = open(inputfile,'rb')                                          
    for i in range(nt):                                              
        for j in range(len(list1)):                                  
            np.fromfile(f, dtype = 'uint32', count = 1)              
            matrix = np.fromfile(f, dtype = 'float64', count = n*m*l)
            list1[j] = np.reshape(matrix, element.shape, order = 'F')     
        if i==0:
            list2[j] = list1[j]
        else:
            list2[j] = np.concatenate((list2[j],list1[j]), axis = 3)             
    f.close()
# Associating names to the elements in the list:
    u, v, w, t, p, bx, by, bz = list2
    return u, v, w, t, p, bx, by, bz



#----------------------------------------------------#
#                   Read XAVERAGES                   #

def rxav(inputfile, m, l, nt):
    '''
    This function is designed for read a XAVERAGES file written in order: empty line, followed by a matrix of order (m,l), followed by a scalar number (time) and more one empty line. 
    This is repeated for 27 element for every step time.
    The order of the 27 matrices is ['tthew', 'uu', 'vv', 'ww', 'oox', 'ooy', 'ooz', 'uu2', 'vv2', 'ww2', 'oox2', 'ooy2', 'ooz2', 'rrwv', 'rrwu', 'rrvu', 'bbx', 
    'bby', 'bbz', 'bbx2', 'bby2', 'bbz2', 'bbzby', 'bbzbx', 'bbxby', 'pp', 'tthe']
    INPUTS: 
    > Inputfile: file to be read
    > m: number of steps in theta direction
    > l: number of steps in radius direction
    > nt: number of time steps
    
    OUTPUT:
    > Matrices in order (m, l, nt)
    '''
    element = np.zeros((m, l))                                         # Creates a element with n lines and m columns
    list1 = [element,]*nt                                              # Creates a list with one element for all times steps
    list2 = [list1,]*27                                                # List for all variables
    f = open(inputfile,'rb')                                           # Read the XAVERAGES file
    m = len(list2)
    for i in range(nt):                                                # Temporal loop
        for j in range(m):                                             # Number of matrices to be read, in XAVERAGES 27. 
            np.fromfile(f, dtype = 'uint32', count = 1)                # Read an empty line    
            matrix = np.fromfile(f, dtype = 'float64', count = m*l)
            list2[j][i] = np.transpose(np.reshape(matrix, element.shape, order = 'F'))  # Fortran order     
            np.fromfile(f, dtype = 'int32', count = 1)                 # Read the time
            np.fromfile(f, dtype = 'int32', count = 1)                 # Read the time
    f.close()
    # Associating names to the elements in the list:
    thew , u , v , w , ox , oy , oz , u2 , v2 , w2 , ox2 , oy2 , oz2 , rwv , rwu , rvu , bx , by , bz , bx2 , by2 , bz2 , bzby , bzbx , bxby , p , the = list2 
    return thew , u , v , w , ox , oy , oz , u2 , v2 , w2 , ox2 , oy2 , oz2 , rwv , rwu , rvu , bx , by , bz , bx2 , by2 , bz2 , bzby , bzbx , bxby , p , the 

#----------------------------------------------------#
#              Read Stratification                   #

def rstrat(inputfile): 
    ''' 
    This function reads the file with stratification and calculates the gamma in function of the radius.
    Is expected that the file has 8 columns, arranged as follows: [r, P, T, rho, cpcv, dltdlp, grad_ad, grad_rad]. 
       
    INPUTS:
    > inputfile: contains the path of the file to be read;
    
    OUTPUT: 
    > r: Radius array;
    > P: Pressure array;
    > T: Temperature array;
    > rho: Density array; 
    > cpcv: Ratio of the specific heat capacity at constant pressure and at constant volumn; 
    > dlntdlnp
    > grad_ad: Adiabatic gradient;
    > grad_rad: Radiative gradient;
    > gamma: Gamma array.
    '''	
    f = open(inputfile, 'r')
    files = np.genfromtxt(f, usecols = (0, 1, 2, 3, 4, 5, 6, 7))
    r = 1e-2*files[:,0]                                           # In meters 
    P = files[:,1]					          
    T = files[:,2]					          # In Kelvin	
    rho = 1e3*files[:,3]					  # In kg/m**3
    cpcv = files[:,4]
    dlntdlnp = files[:,5]
    grad_ad = files[:,6]
    grad_rad = files[:,7]
    gamma = 1./(1. - grad_ad) 
    return r, P, T, rho, grad_ad, grad_rad, gamma									


